import math
import os
from prob import *
import gurobipy as gp
from gurobipy import GRB
import numpy as np

def func_x3(x):
    return x*x*x

def func_sin(x):
    return math.sin(x)

def func_cos(x):
    return math.cos(x)

def func_sin2(x):
    return math.sin(x) * math.sin(x)


def func_expneg(x):
    return math.exp(-x)


def getXt_x3(k, x0, x1):
    m = gp.Model()
    m.setParam('OutputFlag', 0)
    xt = m.addVar(lb=x0, ub=x1, name="xt")
    v = m.addVar(name="v")
    diff = m.addVar(name="diff")
    df = 3 * xt * xt
    m.addConstr(v == df - k)
    m.addGenConstrAbs(diff, v)
    m.setObjective(diff, GRB.MINIMIZE)
    m.Params.NonConvex = 2
    m.optimize()
    # print("getXt_x3 xt.X: ", xt.X)
    return xt.X


def getXt_expneg(k, x0, x1):
    m = gp.Model()
    m.setParam('OutputFlag', 0)
    negxt = m.addVar(lb=-x1, ub=-x0, name="negxt")
    v = m.addVar(name="v")
    negdf = m.addVar(name="negdf")
    diff = m.addVar(name="diff")
    m.addGenConstrExp(negxt, negdf)  # f = e^(-x) ==> -f' = e^(-x)
    m.addConstr(v == -negdf - k)

    m.addGenConstrAbs(diff, v)  # diff = |f' - k|
    m.setObjective(diff, GRB.MINIMIZE)
    m.Params.NonConvex = 2
    m.optimize()
    # print("getXt_expneg xt.X: ", -negxt.X)
    return -negxt.X


def getXt_sin2(k, x0, x1):  # min(|sin(2x) - k|)  s.t. x0 <= x <= x1  ==> e.i. (sinx)^2 =k
    m = gp.Model()
    m.setParam('OutputFlag', 0)
    v = m.addVar(name="v")
    df = m.addVar(name="df")
    diff = m.addVar(name="diff")

    if k > 0:
        xt2 = m.addVar(lb=x0 * 2, ub=x1 * 2, name="xt")
        m.addGenConstrSin(xt2, df)  # f = (sinx)^2 ==> f' = sin(2x)
        m.addConstr(v == df - k)
    else:  # gurobi can only get positive value
        xt2 = m.addVar(lb=x1 * -2, ub=x0 * -2, name="xt")
        m.addGenConstrSin(xt2, df)
        m.addConstr(v == df + k)

    m.addGenConstrAbs(diff, v)
    m.setObjective(diff, GRB.MINIMIZE)
    m.Params.NonConvex = 2
    m.optimize()
    if k > 0:
        # print("getXt_sin2 xt.X: ", xt2.X / 2)
        return xt2.X / 2
    else:
        # print("getXt_sin2 xt.X: ", xt2.X / -2)
        return xt2.X / -2


def getXt_sin(k, x0, x1):
    m = gp.Model()
    m.setParam('OutputFlag', 0)
    v = m.addVar(name="v")
    df = m.addVar(name="df")
    diff = m.addVar(name="diff")

    xt = m.addVar(lb=x0, ub=x1, name="xt")
    m.addGenConstrCos(xt, df)
    m.addConstr(v == df - k)

    m.addGenConstrAbs(diff, v)
    m.setObjective(diff, GRB.MINIMIZE)
    m.Params.NonConvex = 2
    m.optimize()
    return xt.X / 2


def getXt_cos(k, x0, x1):
    m = gp.Model()
    m.setParam('OutputFlag', 0)
    v = m.addVar(name="v")
    negdf = m.addVar(name="negdf")
    diff = m.addVar(name="diff")

    if k > 0:
        xt = m.addVar(lb=-x1, ub=-x0, name="xt")
        m.addGenConstrSin(xt, negdf)  # f = cosx ==> f' = -sinx  // -f' = sinx
        m.addConstr(v == negdf - k)  # |-sinx - k| = |sinx + k|
    else:  # gurobi can only get positive value
        xt = m.addVar(lb=x0, ub=x1, name="xt")
        m.addGenConstrSin(xt, negdf)
        m.addConstr(v == negdf + k)

    m.addGenConstrAbs(diff, v)
    m.setObjective(diff, GRB.MINIMIZE)
    m.Params.NonConvex = 2
    m.optimize()
    if k > 0:
        # print("k>0 xt.X: ", xt.X)
        return xt.X / 2
    else:
        # print("k<0 xt.X: ", xt.X)
        return xt.X / -2


def linear_function_expneg(x0, x1, y0, y1):
    k = (y1 - y0) / (x1 - x0)
    # print("=============linear_function_expneg==============k, x0, x1, y0, y1: ", k, x0, x1, y0, y1)
    b_u = (x1 * y0 - x0 * y1) / (x1 - x0)
    xt = getXt_expneg(k, x0, x1)
    yt = func_expneg(xt)
    b_l = yt - k * xt
    return k, b_u, b_l


def linear_function_x3(x0, x1, y0, y1):
    k = (y1 - y0) / (x1 - x0)
    if x1 <= 0:  # 凹，下界连线，上界切线
        b_l = (x1 * y0 - x0 * y1) / (x1 - x0)
        xt = getXt_x3(k, x0, x1)
        yt = func_x3(xt)
        b_u = yt - k * xt
    else:  # 凸，上界连线，下界切线
        b_u = (x1 * y0 - x0 * y1) / (x1 - x0)
        xt = getXt_x3(k, x0, x1)
        yt = func_x3(xt)
        b_l = yt - k * xt
    return k, b_u, b_l


def linear_function_sin2_convex(x0, x1, y0, y1):  # [-pi/4, pi/4]
    k = (y1 - y0) / (x1 - x0)
    # print("==============linear_function_sin2_convex=============k: ", k)
    b_u = (x1 * y0 - x0 * y1) / (x1 - x0)
    xt = getXt_sin2(k, x0, x1)
    yt = func_sin2(xt)
    b_l = yt - k * xt
    return k, b_u, b_l


def linear_function_sin2_concave(x0, x1, y0, y1):  # [-2, -pi/4] [pi/4, 2]
    k = (y1 - y0) / (x1 - x0)
    b_l = (x1 * y0 - x0 * y1) / (x1 - x0)
    xt = getXt_sin2(k, x0, x1)
    yt = func_sin2(xt)
    b_u = yt - k * xt
    # print("============linear_function_sin2_concave===============k, b_u, b_l: ", k, b_u, b_l)
    return k, b_u, b_l

def linear_function_cos_concave(x0, x1, y0, y1):
    k = (y1 - y0) / (x1 - x0)
    b_l = (x1 * y0 - x0 * y1) / (x1 - x0)
    xt = getXt_cos(k, x0, x1)
    yt = func_cos(xt)
    b_u = yt - k * xt
    return k, b_u, b_l

def linear_function_sin_concave(x0, x1, y0, y1):
    k = (y1 - y0) / (x1 - x0)
    b_l = (x1 * y0 - x0 * y1) / (x1 - x0)
    xt = getXt_sin(k, x0, x1)
    yt = func_sin(xt)
    b_u = yt - k * xt
    return k, b_u, b_l

def linear_function_sin_convex(x0, x1, y0, y1):
    k = (y1 - y0) / (x1 - x0)
    b_u = (x1 * y0 - x0 * y1) / (x1 - x0)
    xt = getXt_sin(k, x0, x1)
    yt = func_sin(xt)
    b_l = yt - k * xt
    return k, b_u, b_l

def piecewise_xy(x0_min, x0_max, x0_num, x1_min, x1_max, x1_num):
    m = gp.Model()
    m.setParam('OutputFlag', 0)
    m.setParam('NonConvex', 2)

    w = m.addMVar((2,), vtype=GRB.CONTINUOUS, name='w', lb=-np.inf, ub=np.inf)

    b_l = m.addVar(lb=-np.inf, ub=np.inf, name="b_l")
    b = m.addVar(lb=0, ub=np.inf, name="b")

    grid_points_x0 = np.linspace(x0_min, x0_max, x0_num)
    grid_points_x1 = np.linspace(x1_min, x1_max, x1_num)

    for x0_val in grid_points_x0:
        for x1_val in grid_points_x1:
            # f = x0_val * x0_val * x1_val  # piecewise_item
            f = x0_val * x1_val  # piecewise_item
            m.addConstr(w[0] * x0_val + w[1] * x1_val + b_l <= f, name=f"lhs_{x0_val}_{x1_val}")
            m.addConstr(f <= w[0] * x0_val + w[1] * x1_val + b_l + b, name=f"rhs_{x0_val}_{x1_val}")

    m.setObjective(b, GRB.MINIMIZE)
    m.optimize()

    print(f"status = {m.status}, in [{x0_min},{x0_max}], [{x1_min},{x1_max}]")
    if m.status == 3:
        m.computeIIS()
        m.write("piecewise.ilp")
        # return None, None, None, None

    if m.status != 3:
        # for v in m.getVars():
        #     print(f"{v.varName} = {v.X}")
        return w[0].X, w[1].X, b_l.X, b.X


def piecewise_xx(x_min, x_max, x_num):
    m = gp.Model()
    m.setParam('OutputFlag', 0)
    m.setParam('NonConvex', 2)

    k = m.addVar(vtype=GRB.CONTINUOUS, name='k', lb=-np.inf, ub=np.inf)
    b_l = m.addVar(lb=-np.inf, ub=np.inf, name="b_l")
    b = m.addVar(lb=0, ub=np.inf, name="b")

    grid_points_x = np.linspace(x_min, x_max, x_num)
    for x_val in grid_points_x:
        expr = gp.QuadExpr(x_val * x_val)  # f = x^2
        expr = gp.QuadExpr(expr * x_val)  # f = x^3
        m.addConstr(k * x_val + b_l <= expr, name=f"lhs_{x_val}")
        m.addConstr(expr <= k * x_val + b_l + b, name=f"rhs_{x_val}")

    m.setObjective(b, GRB.MINIMIZE)
    m.optimize()

    print(f"status = {m.status}, in [{x_min}, {x_max}]")
    if m.status == 3:
        m.computeIIS()
        m.write("piecewise.ilp")
        # return None, None, None, None

    if m.status != 3:
        # for v in m.getVars():
        #     print(f"{v.varName} = {v.X}")
        return k.X, b_l.X, b.X


def generate_piecewise_data_2d(x0_min, x0_max, x1_min, x1_max, x0_len, x1_len, x0_num, x1_num, eps, name):
    # f = x0x1 + linear_expr
    # x0_min = -2, x0_max = 2, x1_min = -2, x1_max = 2, x0_len = 0.1, x1_len = 0.1
    # piece_x0num = 10, piece_x1num = 10, eps = 1e-16
    data_dir = './data/'
    if not os.path.exists(data_dir):
        os.mkdir(data_dir)
    # x0, x1 in range [-2, 2], each axis piece like [-2, -1.9]...
    # x0 in [-2, -1.9] and x1 in [-2, -1.9], and each block size is 0.1*0.1
    # sample the grid points of each block, divided into 10 parts at each axis
    # fw = open(data_dir + "-2_2_0.1_10.txt", "w+")
    fw = open(data_dir + f"{name}.txt", "w+")
    x0_l = x0_min
    x0_r = min(round(x0_l + x0_len, 5), x0_max)
    while x0_l < x0_max:

        k_0, b_l_0, b_0 = piecewise_xx(x0_l, x0_r, x0_num)
        fw.write(str(k_0) + " " + str(b_l_0 - eps) + " " + str(b_0 + 2 * eps) + "\n")

        # x1_l = x1_min
        # x1_r = min(round(x1_l + x1_len, 5), x1_max)
        # while x1_l < x1_max:
        #     w0, w1, b_l, b = piecewise_xy(x0_l, x0_r, x0_num, x1_l, x1_r, x1_num)  # lower_bound and delta_b
        #     # fw.write(str(x0_l) + " " + str(x0_r) + " " + str(x1_l) + " " + str(x1_r) + " " + str(w0) + " " + str(w1) + " " + str(b_l - eps) + " " + str(b + eps) + "\n")
        #     fw.write(str(w0) + " " + str(w1) + " " + str(b_l - eps) + " " + str(b + 2 * eps) + "\n")
        #     x1_l = x1_r
        #     x1_r = min(round(x1_l + x1_len, 5), x1_max)

        x0_l = x0_r
        x0_r = min(round(x0_l + x0_len, 5), x0_max)
    fw.close()


# def generate_piecewise_data_2d(x_min, x_max, x_len, x_num, eps, name):
#     data_dir = './data/'
#     if not os.path.exists(data_dir):
#         os.mkdir(data_dir)
#     # x0, x1 in range [-2, 2], each axis piece like [-2, -1.9]...
#     # x0 in [-2, -1.9] and x1 in [-2, -1.9], and each block size is 0.1*0.1
#     # sample the grid points of each block, divided into 10 parts at each axis
#     # fw = open(data_dir + "-2_2_0.1_10.txt", "w+")
#     fw = open(data_dir + f"{name}.txt", "w+")
#     x_l = x_min
#     x_r = min(x_l + x_len, x_max)
#     while x_l < x_max:
#         k, b_l, b = piecewise_xx(x_l, x_r, x_num)
#         # fw.write(str(x0_l) + " " + str(x0_r) + " " + str(x1_l) + " " + str(x1_r) + " " + str(w0) + " " + str(w1) + " " + str(b_l - eps) + " " + str(b + eps) + "\n")
#         fw.write(str(k) + " " + str(b_l - eps) + " " + str(b + eps) + "\n")
#         x_l = x_r
#         x_r = min(x_l + x_len, x_max)
#     fw.close()


def generate_piecewise_data_2d_in_3d(x0_min, x0_max, x1_min, x1_max, x2_min, x2_max, x0_len, x1_len, x2_len, x0_num, x1_num, x2_num, eps, name):
    # f = x0 * x1 + x0 * x2 + x1 * x2
    # x0_min = -2, x0_max = 2, x1_min = -2, x1_max = 2, x0_len = 0.1, x1_len = 0.1
    # piece_x0num = 10, piece_x1num = 10, eps = 1e-16
    data_dir = './data/'
    if not os.path.exists(data_dir):
        os.mkdir(data_dir)
    # x0, x1 in range [-2, 2], each axis piece like [-2, -1.9]...
    # x0 in [-2, -1.9] and x1 in [-2, -1.9], and each block size is 0.1*0.1
    # sample the grid points of each block, divided into 10 parts at each axis
    # fw = open(data_dir + "-2_2_0.1_10.txt", "w+")
    fw = open(data_dir + f"{name}.txt", "w+")
    x0_l = x0_min
    x0_r = min(x0_l + x0_len, x0_max)
    while x0_l < x0_max:
        x0_k, x0_b_l, x0_b = piecewise_xx(x0_l, x0_r, x0_num)
        x1_l = x1_min
        x1_r = min(x1_l + x1_len, x1_max)
        while x1_l < x1_max:
            x1_k, x1_b_l, x1_b = piecewise_xx(x1_l, x1_r, x1_num)
            x01_w0, x01_w1, x01_b_l, x01_b = piecewise_xy(x0_l, x0_r, x0_num, x1_l, x1_r, x1_num)
            x2_l = x2_min
            x2_r = min(x2_l + x2_len, x2_max)
            while x2_l < x2_max:
                x2_k, x2_b_l, x2_b = piecewise_xx(x2_l, x2_r, x2_num)
                x02_w0, x02_w1, x02_b_l, x02_b = piecewise_xy(x0_l, x0_r, x0_num, x2_l, x2_r, x2_num)
                x12_w0, x12_w1, x12_b_l, x12_b = piecewise_xy(x1_l, x1_r, x1_num, x2_l, x2_r, x2_num)
                fw.write(str(x0_l) + " " + str(x0_r) + " " + str(x1_l) + " " + str(x1_r) + " " + str(x2_l) + " " + str(x2_r) + " ")
                fw.write(str(x01_w0) + " " + str(x01_w1) + " " + str(x01_b_l - eps) + " " + str(x01_b + eps) + " ")
                fw.write(str(x02_w0) + " " + str(x02_w1) + " " + str(x02_b_l - eps) + " " + str(x02_b + eps) + " ")
                fw.write(str(x12_w0) + " " + str(x12_w1) + " " + str(x12_b_l - eps) + " " + str(x12_b + eps) + " ")
                fw.write(str(x0_k) + " " + str(x0_b_l - eps) + " " + str(x0_b + eps) + " ")
                fw.write(str(x1_k) + " " + str(x1_b_l - eps) + " " + str(x1_b + eps) + " ")
                fw.write(str(x2_k) + " " + str(x2_b_l - eps) + " " + str(x2_b + eps) + "\n")
                x2_l = x2_r
                x2_r = min(x2_l + x2_len, x2_max)
            x1_l = x1_r
            x1_r = min(x1_l + x1_len, x1_max)
        x0_l = x0_r
        x0_r = min(x0_l + x0_len, x0_max)
    fw.close()


if __name__ == "__main__":
    begin = -4
    end = 4
    len = 1
    nums = 20
    eps = 1e-10
    name = f"03_{begin}_{end}_{len}_{nums}"
    # name = f"{begin}_{end}_{len}_{nums}"
    generate_piecewise_data_2d(begin, end, begin, end, len, len, nums, nums, eps, name)
    # generate_piecewise_data_2d(-3, 3, 0.05, 50, 0.001, "-3_3_0.05_20_xx")
    # generate_piecewise_data_2d_in_3d(-2, 2, -2, 2, -2, 2, 0.5, 0.5, 0.5, 20,20, 20, 0.001, "-2_2_0.5_20_3d")
    # generate_piecewise_data_2d_in_3d(-2, 2, -2, 2, -2, 2, 2, 2, 2, 20,20, 20, 0.001, "-2_2_2_20_3d")
    # generate_piecewise_data(-2, -1.9, 1.6, 1.7, 0.1, 0.1, 10,10, 0.1, "10points_0.1")
    # k, b_l, b = piecewise_xx(-5, -4, 10)
    # print(f"k={k}, b_l={b_l}, b={b}")
